#include "slam_ros/mute_grid.h"

template<class FloatType>
int MapModifier::WritePbStream(string& read_file,
                   string& write_file,std::map<int, Eigen::Vector3d> point_loc_add_list,std::map<int, Eigen::Vector3d> point_loc_delete_list)
{
  using Vector = Eigen::Matrix<FloatType, 3, 1>;
  using Quaternion = Eigen::Quaternion<FloatType>;
  using AngleAxis = Eigen::AngleAxis<FloatType>;

  vector <slam::transform::proto::Rigid3d> submap_global_pose;
  ::slam::io::ProtoStreamReader reader(read_file);
  ::slam::io::ProtoStreamDeserializer deserializer(&reader);
  ::slam::mapping::proto::SerializedData proto_;
  ::slam::io::ProtoStreamReader reader_(read_file);
  ::slam::io::ProtoStreamWriter writer(write_file);
  ::slam::mapping::proto::SerializationHeader header;
  reader_.ReadProto(&header);
  writer.WriteProto(header);

  const mapping::proto::PoseGraph pose_graph = deserializer.pose_graph();
  const map<SerializedData::DataCase, string> data_case_to_name = {
      {SerializedData::kSubmap, "submap"},
      {SerializedData::kNode, "node"},
      {SerializedData::kTrajectoryData, "trajectory_data"},
      {SerializedData::kImuData, "imu_data"},
      {SerializedData::kOdometryData, "odometry_data"},
      {SerializedData::kFixedFramePoseData, "fixed_frame_pose_data"},
      {SerializedData::kLandmarkData, "landmark_data"},
  };

  std::map<std::string, int> data_counts = {
      {"submap_2d", 0},
      {"submap_2d_grid", 0},
      {"submap_3d", 0},
      {"submap_3d_high_resolution_hybrid_grid", 0},
  };

  int num = 0;
  int n = 0;

  std::map<int, Eigen::Vector3d> mute_pose_add,mute_pose_del;
  mute_pose_add = point_loc_add_list;
  mute_pose_del = point_loc_delete_list;

  std::map<int, Eigen::Vector3d>::iterator mute_pose_iter;
  Eigen::Vector3d GlobalPose;
  while(reader_.ReadProto(&proto_))
  {
    if (proto_.has_pose_graph()){
      for(int k = 0; k < proto_.pose_graph().trajectory(0).submap_size(); k++){
        submap_global_pose.push_back(proto_.pose_graph().trajectory(0).submap(k).pose());
      }
    }
    if (proto_.submap().has_submap_2d()) {
      double submap_max_x_, submap_min_x_, submap_max_y_, submap_min_y_, dif_global_local_x, dif_global_local_y;
      Eigen::Quaterniond dif_global_local_quan;
      dif_global_local_quan.w() = submap_global_pose[num].rotation().w();
      dif_global_local_quan.z() = submap_global_pose[num].rotation().z();
      double m11 = +1.0 - 2.0 * (dif_global_local_quan.z() * dif_global_local_quan.z());
      double m12 = -2.0 * (dif_global_local_quan.w() * dif_global_local_quan.z());
      double m21 = 2.0 * (dif_global_local_quan.w() * dif_global_local_quan.z());
      double m22 = +1.0 - 2.0 * (dif_global_local_quan.z() * dif_global_local_quan.z());
      Eigen::Matrix3d GlobalToLocal; 
      GlobalToLocal << m11, m12, proto_.submap().submap_2d().local_pose().translation().x(),
                        m21, m22, proto_.submap().submap_2d().local_pose().translation().y(),
                        0, 0, 1;
      submap_max_x_ = proto_.submap().submap_2d().grid().limits().max().x();
      submap_max_y_ = proto_.submap().submap_2d().grid().limits().max().y();
      submap_min_x_ = submap_max_x_ - proto_.submap().submap_2d().grid().limits().resolution() * proto_.submap().submap_2d().grid().limits().cell_limits().num_y_cells();
      submap_min_y_ = submap_max_y_ - proto_.submap().submap_2d().grid().limits().resolution() * proto_.submap().submap_2d().grid().limits().cell_limits().num_x_cells();
      for (mute_pose_iter = mute_pose_add.begin(); mute_pose_iter != mute_pose_add.end(); mute_pose_iter ++){
        GlobalPose[0] = mute_pose_iter->second[0];
        GlobalPose[1] = mute_pose_iter->second[1];
        GlobalPose[2] = 1;
        Eigen::Vector3d change_pose_ = GlobalToLocal.inverse() * GlobalPose; 
        double x_dif = submap_max_x_ - (change_pose_[0] - submap_global_pose[num].translation().x() + 2 * proto_.submap().submap_2d().local_pose().translation().x());
        double y_dif = submap_max_y_ - (change_pose_[1] - submap_global_pose[num].translation().y() + 2 * proto_.submap().submap_2d().local_pose().translation().y());
        if (x_dif >= 0 && y_dif >= 0){
          int x_dif_index = x_dif / proto_.submap().submap_2d().grid().limits().resolution();
          int y_dif_index = y_dif / proto_.submap().submap_2d().grid().limits().resolution();
          int index = x_dif_index * proto_.submap().submap_2d().grid().limits().cell_limits().num_x_cells() + y_dif_index;
          if (x_dif_index < proto_.submap().submap_2d().grid().limits().cell_limits().num_y_cells() && y_dif_index < proto_.submap().submap_2d().grid().limits().cell_limits().num_x_cells()){
            proto_.mutable_submap()->mutable_submap_2d()->mutable_grid()->mutable_cells()->Set(index, 100);
          }
        }
      }
       for (mute_pose_iter = mute_pose_del.begin(); mute_pose_iter != mute_pose_del.end(); mute_pose_iter ++){
        GlobalPose[0] = mute_pose_iter->second[0];
        GlobalPose[1] = mute_pose_iter->second[1];
        GlobalPose[2] = 1;
        Eigen::Vector3d change_pose_ = GlobalToLocal.inverse() * GlobalPose;
        double x_dif = submap_max_x_ - (change_pose_[0] - submap_global_pose[num].translation().x() + 2 * proto_.submap().submap_2d().local_pose().translation().x());
        double y_dif = submap_max_y_ - (change_pose_[1] - submap_global_pose[num].translation().y() + 2 * proto_.submap().submap_2d().local_pose().translation().y());
        if (x_dif >= 0 && y_dif >= 0){
          int x_dif_index = x_dif / proto_.submap().submap_2d().grid().limits().resolution();
          int y_dif_index = y_dif / proto_.submap().submap_2d().grid().limits().resolution();
          int index = x_dif_index * proto_.submap().submap_2d().grid().limits().cell_limits().num_x_cells() + y_dif_index;
          if (x_dif_index < proto_.submap().submap_2d().grid().limits().cell_limits().num_y_cells() && y_dif_index < proto_.submap().submap_2d().grid().limits().cell_limits().num_x_cells()){
            proto_.mutable_submap()->mutable_submap_2d()->mutable_grid()->mutable_cells()->Set(index, -1);
          }
        }
      }
     
      num ++;
      RCLCPP_INFO(this->get_logger(), "%d", num);
    }
  writer.WriteProto(proto_);
  n++;
  }

  for (const auto& entry : data_counts) {
    RCLCPP_INFO(this->get_logger(), "SerializedData package contains #%s: %d", 
                entry.first.c_str(), entry.second);
  }

  writer.Close();
  return 0;  
}

bool MapModifier::init()
{
  string test_file = "/home/carto/map/1682045244273.pbstream";
  string write_file = "/home/carto/map/1682045244273_2.pbstream";
  string map_yaml_path = "/home/carto/map/1682045244273.yaml";

  YAML::Node loaded = YAML::LoadFile(map_yaml_path);
  unsigned int cells_size_x = 672;
  unsigned int cells_size_y = 619;

  double resolution  = loaded["resolution"].as<double>();
  double origin_x = loaded["origin"].as<vector<double>>().at(0);
  double origin_y = loaded["origin"].as<vector<double>>().at(1);

  vector<vector<loc_w>> area_delete_list;
  vector<vector<loc_w>> area_add_list;
  vector<loc_w> area1,area2;  

  loc_w point1,point2,point3,point4;
  point1.x = 0.0;
  point1.x = 0.0;
  point2.x = 0.0;
  point2.x = 0.0;
  point3.x = 0.0;
  point3.x = 0.0;
  point4.x = 0.0;
  point4.x = 0.0;
  // point1.x= loaded["area1"]["point1"].as<vector<double>>().at(0);
  // point1.y= loaded["area1"]["point1"].as<vector<double>>().at(1);
  area1.push_back(point1);
  // point2.x= loaded["area1"]["point2"].as<vector<double>>().at(0);
  // point2.y= loaded["area1"]["point2"].as<vector<double>>().at(1);
  area1.push_back(point2);
  // point3.x= loaded["area1"]["point3"].as<vector<double>>().at(0);
  // point3.y= loaded["area1"]["point3"].as<vector<double>>().at(1);
  area1.push_back(point3);
  // point4.x= loaded["area1"]["point4"].as<vector<double>>().at(0);
  // point4.y= loaded["area1"]["point4"].as<vector<double>>().at(1);
  area1.push_back(point4);
  area_add_list.push_back(area1);

  point1.x = -6.0;
  point1.x = 0.0;
  point2.x = -6.0;
  point2.x = 10.0;
  point3.x = -15.0;
  point3.x = 10.0;
  point4.x = -15.0;
  point4.x = 0.0;
  // point1.x= loaded["area2"]["point1"].as<vector<double>>().at(0);
  // point1.y= loaded["area2"]["point1"].as<vector<double>>().at(1);
  area2.push_back(point1);
  // point2.x= loaded["area2"]["point2"].as<vector<double>>().at(0);
  // point2.y= loaded["area2"]["point2"].as<vector<double>>().at(1);
  area2.push_back(point2);
  // point3.x= loaded["area2"]["point3"].as<vector<double>>().at(0);
  // point3.y= loaded["area2"]["point3"].as<vector<double>>().at(1);
  area2.push_back(point3);
  // point4.x= loaded["area2"]["point4"].as<vector<double>>().at(0);
  // point4.y= loaded["area2"]["point4"].as<vector<double>>().at(1);
  area2.push_back(point4);
  area_delete_list.push_back(area2);

  this->costMap = Costmap2D(cells_size_x, cells_size_y, resolution,
                     origin_x, origin_y);

  std::map<int, Eigen::Vector3d> point_loc_add_list = this->CalculatePoints(area_add_list);
  std::map<int, Eigen::Vector3d> point_loc_delete_list = this->CalculatePoints(area_delete_list);
  this->WritePbStream<double>(test_file, write_file,point_loc_add_list,point_loc_delete_list);
  return true;
}

std::map<int, Eigen::Vector3d> MapModifier::CalculatePoints(vector<vector<loc_w>> polygonList)
{
  std::map<int, Eigen::Vector3d> world_loc_list;
  for (vector<loc_w> polygon : polygonList){
      vector<MapLocation> map_loc_list;
      vector<MapLocation> map_loc_inner_list;      
      for(loc_w point_w : polygon){
        MapLocation map_loc;
        this->costMap.worldToMap(point_w.x,point_w.y,map_loc.x,map_loc.y);
        map_loc_list.push_back(map_loc);
      }
      this->costMap.convexFillCells(map_loc_list,map_loc_inner_list);
      for(MapLocation point_m : map_loc_inner_list){
        loc_w world_loc;
        this->costMap.mapToWorld(point_m.x,point_m.y,world_loc.x,world_loc.y);
        int len = world_loc_list.size();
        std::map<int, Eigen::Vector3d> world_loc_;
        world_loc_list[len+1] = {world_loc.x,world_loc.y,1};
      }
   }
  return world_loc_list;
}

MapModifier::MapModifier(const rclcpp::NodeOptions & options)
: Node("map_modifier_node", options)
{
    if(!init()) {
        RCLCPP_ERROR(this->get_logger(), "Failed to initialize costmap");
        throw std::runtime_error("Costmap initialization failed");
    }
}

MapModifier::~MapModifier()
{
    
}